import os
import time

import yaml

from web_live2.page_object.base_page import BasePage



class LoginPage(BasePage):

    _BASE_URL = 'https://work.weixin.qq.com/wework_admin/frame'

    def get_cookies(self):
        time.sleep(20)
        cookies = self.driver.get_cookies()
        with open("../data/cookie.yaml", "w") as f:
            yaml.safe_dump(cookies, f)
        self.do_get(self._BASE_URL)

    def add_cookies(self):
        cookies = yaml.safe_load(open("../data/cookie.yaml"))
        for cookie in cookies:
            self.driver.add_cookie(cookie)

        self.do_get(self._BASE_URL)

    def login(self):

        if os.path.exists("../data/cookie.yaml"):
            self.add_cookies()

        else:
            self.get_cookies()

        from web_live2.page_object.home_page import HomePage
        return HomePage(self.driver)